{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bedrock Agent Integration with Bedrock AgentCore Gateway\n",
    "\n",
    "This notebook demonstrates end-to-end integration of Amazon Bedrock Agents with Bedrock AgentCore Gateway for a Fruit Stand Backend API that is built on AWS Lambda and DynamoDB\n",
    "\n",
    "## Architecture Overview\n",
    "```\n",
    "User Query → Bedrock Agent → Bridge Lambda → Bedrock AgentCore Gateway → Target Lambda → DynamoDB\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Install Dependencies & Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install boto3 \n",
    "!pip3 install botocore--quiet\n",
    "!pip3 install boto3 --quiet\n",
    "\n",
    "import boto3\n",
    "import json\n",
    "import zipfile\n",
    "import os\n",
    "import urllib.request\n",
    "import urllib.parse\n",
    "import time\n",
    "\n",
    "print(\"✅ Dependencies installed successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create IAM Role with necessary permissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_iam_role():\n",
    "    \"\"\"Create IAM role that can be used for Bedrock Core Gateway, Bedrock Agent, Lambda\"\"\"\n",
    "    iam = boto3.client('iam')\n",
    "    \n",
    "    trust_policy = {\n",
    "        \"Version\": \"2012-10-17\",\n",
    "        \"Statement\": [\n",
    "            {\n",
    "                \"Effect\": \"Allow\",\n",
    "                \"Principal\": {\n",
    "                    \"Service\": \"lambda.amazonaws.com\"\n",
    "                },\n",
    "                \"Action\": \"sts:AssumeRole\"\n",
    "            },\n",
    "            {\n",
    "                \"Effect\": \"Allow\",\n",
    "                \"Principal\": {\n",
    "                    \"Service\": \"bedrock.amazonaws.com\"\n",
    "                },\n",
    "                \"Action\": \"sts:AssumeRole\"\n",
    "            },\n",
    "            {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Principal\": {\n",
    "                \"Service\": \"bedrock-agentcore.amazonaws.com\"\n",
    "            },\n",
    "            \"Action\": \"sts:AssumeRole\"\n",
    "        }\n",
    "        ]\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        role_response = iam.create_role(\n",
    "            RoleName='BedrockAgentCoreGatewayLambdaRole',\n",
    "            AssumeRolePolicyDocument=json.dumps(trust_policy),\n",
    "            Description='Role for Bedrock Gateway with DynamoDB access'\n",
    "        )\n",
    "        \n",
    "        # Create and attach inline policy\n",
    "        permissions_policy = {\n",
    "            \"Version\": \"2012-10-17\",\n",
    "            \"Statement\": [\n",
    "                {\n",
    "                    \"Effect\": \"Allow\",\n",
    "                    \"Action\": [\n",
    "                        \"bedrock-agentcore:*\",\n",
    "                        \"iam:PassRole\",\n",
    "                        \"lambda:InvokeFunction\"\n",
    "                    ],\n",
    "                    \"Resource\": \"*\"\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "        \n",
    "        iam.put_role_policy(\n",
    "            RoleName='BedrockAgentCoreGatewayLambdaRole',\n",
    "            PolicyName='BedrockAgentCoreInlinePolicy',\n",
    "            PolicyDocument=json.dumps(permissions_policy)\n",
    "        )\n",
    "        \n",
    "        # Attach managed policies\n",
    "        iam.attach_role_policy(\n",
    "            RoleName='BedrockAgentCoreGatewayLambdaRole',\n",
    "            PolicyArn='arn:aws:iam::aws:policy/AdministratorAccess'\n",
    "        )\n",
    "        \n",
    "        role_arn = role_response['Role']['Arn']\n",
    "        print(f\"✅ Created IAM role: {role_arn}\")\n",
    "        return role_arn\n",
    "        \n",
    "    except Exception as e:\n",
    "        if 'already exists' in str(e):\n",
    "            # Ensure existing role has all required permissions\n",
    "            try:\n",
    "                iam.attach_role_policy(\n",
    "                    RoleName='BedrockAgentCoreGatewayLambdaRole',\n",
    "                    PolicyArn='arn:aws:iam::aws:policy/AdministratorAccess'\n",
    "                )\n",
    "            except:\n",
    "                pass  # Policy may already be attached\n",
    "            \n",
    "            account_id = boto3.client('sts').get_caller_identity()['Account']\n",
    "            role_arn = f\"arn:aws:iam::{account_id}:role/BedrockAgentCoreGatewayLambdaRole\"\n",
    "            print(f\"✅ Using existing IAM role: {role_arn}\")\n",
    "            return role_arn\n",
    "        else:\n",
    "            print(f\"❌ Error creating role: {e}\")\n",
    "            return None\n",
    "\n",
    "gateway_role_arn = create_iam_role()\n",
    "print(f\"Gateway Role ARN: {gateway_role_arn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Create DynamoDB Tables\n",
    "\n",
    "Create DynamoDB tables for orders and inventory tracking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dynamodb_tables():\n",
    "    \"\"\"Create DynamoDB tables for orders and inventory\"\"\"\n",
    "    dynamodb = boto3.client('dynamodb', region_name='us-east-1')\n",
    "    \n",
    "    # Create Orders table\n",
    "    try:\n",
    "        orders_table = dynamodb.create_table(\n",
    "            TableName='FruitOrders',\n",
    "            KeySchema=[\n",
    "                {'AttributeName': 'order_id', 'KeyType': 'HASH'}\n",
    "            ],\n",
    "            AttributeDefinitions=[\n",
    "                {'AttributeName': 'order_id', 'AttributeType': 'S'}\n",
    "            ],\n",
    "            BillingMode='PAY_PER_REQUEST'\n",
    "        )\n",
    "        print(\"✅ Created Orders table\")\n",
    "    except Exception as e:\n",
    "        if 'already exists' in str(e):\n",
    "            print(\"✅ Orders table already exists\")\n",
    "        else:\n",
    "            print(f\"❌ Error creating Orders table: {e}\")\n",
    "    \n",
    "    # Create Inventory table\n",
    "    try:\n",
    "        inventory_table = dynamodb.create_table(\n",
    "            TableName='FruitInventory',\n",
    "            KeySchema=[\n",
    "                {'AttributeName': 'fruit_name', 'KeyType': 'HASH'}\n",
    "            ],\n",
    "            AttributeDefinitions=[\n",
    "                {'AttributeName': 'fruit_name', 'AttributeType': 'S'}\n",
    "            ],\n",
    "            BillingMode='PAY_PER_REQUEST'\n",
    "        )\n",
    "        print(\"✅ Created Inventory table\")\n",
    "        \n",
    "        # Initialize inventory data\n",
    "        # Wait for table to be active\n",
    "        waiter = dynamodb.get_waiter(\"table_exists\")\n",
    "        waiter.wait(TableName=\"FruitInventory\", WaiterConfig={\"Delay\": 2, \"MaxAttempts\": 30})\n",
    "        print(\"✅ Table is active\")\n",
    "        \n",
    "        \n",
    "        fruits_data = [\n",
    "            {'fruit_name': 'apple', 'price': 1.20, 'unit': 'each', 'stock': 100},\n",
    "            {'fruit_name': 'banana', 'price': 0.50, 'unit': 'each', 'stock': 150},\n",
    "            {'fruit_name': 'orange', 'price': 0.75, 'unit': 'each', 'stock': 80},\n",
    "            {'fruit_name': 'strawberry', 'price': 3.99, 'unit': 'pound', 'stock': 25},\n",
    "            {'fruit_name': 'blueberry', 'price': 4.50, 'unit': 'pint', 'stock': 30}\n",
    "        ]\n",
    "        \n",
    "        for fruit in fruits_data:\n",
    "            dynamodb.put_item(\n",
    "                TableName='FruitInventory',\n",
    "                Item={\n",
    "                    'fruit_name': {'S': fruit['fruit_name']},\n",
    "                    'price': {'N': str(fruit['price'])},\n",
    "                    'unit': {'S': fruit['unit']},\n",
    "                    'stock': {'N': str(fruit['stock'])}\n",
    "                }\n",
    "            )\n",
    "        print(\"✅ Initialized inventory data\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        if 'already exists' in str(e):\n",
    "            print(\"✅ Inventory table already exists\")\n",
    "        else:\n",
    "            print(f\"❌ Error creating Inventory table: {e}\")\n",
    "\n",
    "# Create DynamoDB tables\n",
    "create_dynamodb_tables()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Setup Cognito Authentication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_cognito_auth():\n",
    "    \"\"\"Set up Cognito authentication for Bedrock AgentCore Gateway\"\"\"\n",
    "    cognito = boto3.client('cognito-idp', region_name='us-east-1')\n",
    "    \n",
    "    try:\n",
    "        user_pool_response = cognito.create_user_pool(\n",
    "            PoolName='bedrock-agentcore-gateway-dynamo-pool'\n",
    "        )\n",
    "        user_pool_id = user_pool_response['UserPool']['Id']\n",
    "        \n",
    "        import random\n",
    "        import string\n",
    "        domain_name = f\"bedrock-agentcore-dynamo-{''.join(random.choices(string.ascii_lowercase + string.digits, k=8))}\"\n",
    "        \n",
    "        try:\n",
    "            cognito.create_user_pool_domain(\n",
    "                Domain=domain_name,\n",
    "                UserPoolId=user_pool_id\n",
    "            )\n",
    "        except:\n",
    "            pass\n",
    "        \n",
    "        cognito.create_resource_server(\n",
    "            UserPoolId=user_pool_id,\n",
    "            Identifier='bedrock-agentcore-server',\n",
    "            Name='BedrockCoreGatewayServer',\n",
    "            Scopes=[\n",
    "                {'ScopeName': 'read', 'ScopeDescription': 'Read access'},\n",
    "                {'ScopeName': 'write', 'ScopeDescription': 'Write access'}\n",
    "            ]\n",
    "        )\n",
    "        \n",
    "        client_response = cognito.create_user_pool_client(\n",
    "            UserPoolId=user_pool_id,\n",
    "            ClientName='bedrock-agentcore-client',\n",
    "            GenerateSecret=True,\n",
    "            AllowedOAuthFlows=['client_credentials'],\n",
    "            AllowedOAuthScopes=['bedrock-agentcore-server/read', 'bedrock-agentcore-server/write'],\n",
    "            AllowedOAuthFlowsUserPoolClient=True,\n",
    "            SupportedIdentityProviders=['COGNITO']\n",
    "        )\n",
    "        \n",
    "        client_id = client_response['UserPoolClient']['ClientId']\n",
    "        client_secret = client_response['UserPoolClient']['ClientSecret']\n",
    "        \n",
    "        auth_config = {\n",
    "            'user_pool_id': user_pool_id,\n",
    "            'client_id': client_id,\n",
    "            'client_secret': client_secret,\n",
    "            'discovery_url': f\"https://cognito-idp.us-east-1.amazonaws.com/{user_pool_id}/.well-known/openid-configuration\",\n",
    "            'token_endpoint': f\"https://{domain_name}.auth.us-east-1.amazoncognito.com/oauth2/token\",\n",
    "            'domain_name': domain_name\n",
    "        }\n",
    "        \n",
    "        print(f\"✅ Created Cognito configuration\")\n",
    "        return auth_config\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error setting up Cognito: {e}\")\n",
    "\n",
    "auth_config = setup_cognito_auth()\n",
    "print(f\"Auth Configuration: {json.dumps(auth_config, indent=2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Create Backend API Lambda Function with DynamoDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enhanced Lambda function code with DynamoDB integration\n",
    "lambda_code = '''\n",
    "import json\n",
    "import uuid\n",
    "import boto3\n",
    "from datetime import datetime\n",
    "from decimal import Decimal\n",
    "\n",
    "dynamodb = boto3.resource('dynamodb', region_name='us-east-1')\n",
    "orders_table = dynamodb.Table('FruitOrders')\n",
    "inventory_table = dynamodb.Table('FruitInventory')\n",
    "\n",
    "def determine_tool_from_arguments(event):\n",
    "    if not event:\n",
    "        return \"list_fruits_tool\"\n",
    "    if \"customer_name\" in event or \"items\" in event:\n",
    "        return \"create_order_tool\"\n",
    "    elif \"order_id\" in event:\n",
    "        return \"get_order_tool\"\n",
    "    return \"list_fruits_tool\"\n",
    "\n",
    "def list_fruits_tool(arguments):\n",
    "    try:\n",
    "        response = inventory_table.scan()\n",
    "        fruits = []\n",
    "        for item in response['Items']:\n",
    "            fruits.append({\n",
    "                \"name\": item['fruit_name'],\n",
    "                \"price\": float(item['price']),\n",
    "                \"unit\": item['unit'],\n",
    "                \"stock\": int(item['stock'])\n",
    "            })\n",
    "        return {\"fruits\": fruits, \"status\": \"success\"}\n",
    "    except Exception as e:\n",
    "        return {\"error\": str(e), \"status\": \"error\"}\n",
    "\n",
    "def create_order_tool(arguments):\n",
    "    try:\n",
    "        items = arguments.get('items', [])\n",
    "        customer_name = arguments.get('customer_name', 'Anonymous')\n",
    "        \n",
    "        if not items:\n",
    "            return {\"error\": \"No items provided\", \"status\": \"error\"}\n",
    "        \n",
    "        order_id = str(uuid.uuid4())[:8]\n",
    "        total_cost = Decimal('0')\n",
    "        order_items = []\n",
    "        \n",
    "        for item in items:\n",
    "            fruit_name = item.get('name', '').lower()\n",
    "            quantity = item.get('quantity', 1)\n",
    "            \n",
    "            try:\n",
    "                fruit_response = inventory_table.get_item(\n",
    "                    Key={'fruit_name': fruit_name}\n",
    "                )\n",
    "                if 'Item' not in fruit_response:\n",
    "                    return {\"error\": f\"Fruit '{fruit_name}' not found\", \"status\": \"error\"}\n",
    "                \n",
    "                fruit = fruit_response['Item']\n",
    "                if int(fruit['stock']) < quantity:\n",
    "                    return {\"error\": f\"Insufficient stock for {fruit_name}. Available: {fruit['stock']}\", \"status\": \"error\"}\n",
    "                \n",
    "                unit_price = Decimal(str(fruit['price']))\n",
    "                item_cost = unit_price * Decimal(str(quantity))\n",
    "                total_cost += item_cost\n",
    "                \n",
    "                order_items.append({\n",
    "                    \"name\": fruit['fruit_name'],\n",
    "                    \"quantity\": quantity,\n",
    "                    \"unit_price\": unit_price,\n",
    "                    \"unit\": fruit['unit'],\n",
    "                    \"total_price\": item_cost\n",
    "                })\n",
    "                \n",
    "                inventory_table.update_item(\n",
    "                    Key={'fruit_name': fruit_name},\n",
    "                    UpdateExpression='SET stock = stock - :qty',\n",
    "                    ExpressionAttributeValues={':qty': quantity}\n",
    "                )\n",
    "                \n",
    "            except Exception as e:\n",
    "                return {\"error\": f\"Error processing {fruit_name}: {str(e)}\", \"status\": \"error\"}\n",
    "        \n",
    "        orders_table.put_item(Item={\n",
    "            'order_id': order_id,\n",
    "            'customer_name': customer_name,\n",
    "            'items': order_items,\n",
    "            'total_cost': total_cost,\n",
    "            'status': 'pending',\n",
    "            'created_at': datetime.now().isoformat()\n",
    "        })\n",
    "        \n",
    "        order = {\n",
    "            \"order_id\": order_id,\n",
    "            \"customer_name\": customer_name,\n",
    "            \"items\": [\n",
    "                {\n",
    "                    \"name\": item[\"name\"],\n",
    "                    \"quantity\": item[\"quantity\"],\n",
    "                    \"unit_price\": float(item[\"unit_price\"]),\n",
    "                    \"unit\": item[\"unit\"],\n",
    "                    \"total_price\": float(item[\"total_price\"])\n",
    "                }\n",
    "                for item in order_items\n",
    "            ],\n",
    "            \"total_cost\": float(total_cost),\n",
    "            \"status\": \"pending\",\n",
    "            \"created_at\": datetime.now().isoformat()\n",
    "        }\n",
    "        \n",
    "        return {\"order\": order, \"status\": \"success\"}\n",
    "        \n",
    "    except Exception as e:\n",
    "        return {\"error\": str(e), \"status\": \"error\"}\n",
    "\n",
    "def get_order_tool(arguments):\n",
    "    try:\n",
    "        order_id = arguments.get('order_id')\n",
    "        if not order_id:\n",
    "            return {\"error\": \"Order ID is required\", \"status\": \"error\"}\n",
    "        \n",
    "        response = orders_table.get_item(Key={'order_id': order_id})\n",
    "        if 'Item' not in response:\n",
    "            return {\"error\": f\"Order {order_id} not found\", \"status\": \"error\"}\n",
    "        \n",
    "        order = response['Item']\n",
    "        order['total_cost'] = float(order['total_cost'])\n",
    "        for item in order.get('items', []):\n",
    "            if 'unit_price' in item:\n",
    "                item['unit_price'] = float(item['unit_price'])\n",
    "            if 'total_price' in item:\n",
    "                item['total_price'] = float(item['total_price'])\n",
    "        \n",
    "        return {\"order\": order, \"status\": \"success\"}\n",
    "        \n",
    "    except Exception as e:\n",
    "        return {\"error\": str(e), \"status\": \"error\"}\n",
    "\n",
    "def lambda_handler(event, context):\n",
    "    try:\n",
    "        tool_name = determine_tool_from_arguments(event)\n",
    "        \n",
    "        tools = {\n",
    "            'list_fruits_tool': list_fruits_tool,\n",
    "            'create_order_tool': create_order_tool,\n",
    "            'get_order_tool': get_order_tool\n",
    "        }\n",
    "        \n",
    "        tool_function = tools.get(tool_name)\n",
    "        if not tool_function:\n",
    "            return {\"error\": f\"Unknown tool: {tool_name}\", \"status\": \"error\"}\n",
    "        \n",
    "        return tool_function(event)\n",
    "        \n",
    "    except Exception as e:\n",
    "        return {\"error\": str(e), \"status\": \"error\"}\n",
    "'''\n",
    "\n",
    "def create_enhanced_lambda():\n",
    "    lambda_client = boto3.client('lambda', region_name='us-east-1')\n",
    "    \n",
    "    with open('enhanced_lambda_function.py', 'w') as f:\n",
    "        f.write(lambda_code)\n",
    "    \n",
    "    with zipfile.ZipFile('enhanced_fruit_orders_lambda.zip', 'w') as zip_file:\n",
    "        zip_file.write('enhanced_lambda_function.py', 'lambda_function.py')\n",
    "    \n",
    "    try:\n",
    "        with open('enhanced_fruit_orders_lambda.zip', 'rb') as zip_file:\n",
    "            response = lambda_client.create_function(\n",
    "                FunctionName='enhanced-fruit-order-dynamo',\n",
    "                Runtime='python3.13',\n",
    "                Role=gateway_role_arn,\n",
    "                Handler='lambda_function.lambda_handler',\n",
    "                Code={'ZipFile': zip_file.read()},\n",
    "                Description='Enhanced fruit ordering system with DynamoDB',\n",
    "                Timeout=30\n",
    "            )\n",
    "        \n",
    "        print(f\"✅ Created Enhanced Lambda: {response['FunctionArn']}\")\n",
    "        return response['FunctionArn']\n",
    "        \n",
    "    except Exception as e:\n",
    "        if 'already exists' in str(e):\n",
    "            account_id = boto3.client('sts').get_caller_identity()['Account']\n",
    "            lambda_arn = f\"arn:aws:lambda:us-east-1:{account_id}:function:enhanced-fruit-order-dynamo\"\n",
    "            print(f\"✅ Lambda function already exists: {lambda_arn}\")\n",
    "            return lambda_arn\n",
    "        else:\n",
    "            print(f\"❌ Error: {e}\")\n",
    "            return None\n",
    "    finally:\n",
    "        if os.path.exists('enhanced_lambda_function.py'):\n",
    "            os.remove('enhanced_lambda_function.py')\n",
    "        if os.path.exists('enhanced_fruit_orders_lambda.zip'):\n",
    "            os.remove('enhanced_fruit_orders_lambda.zip')\n",
    "\n",
    "enhanced_lambda_arn = create_enhanced_lambda()\n",
    "print(f\"Enhanced Lambda ARN: {enhanced_lambda_arn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Create Bedrock AgentCore Gateway"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_bedrock_agentcore_gateway(role_arn, auth_config):\n",
    "    \"\"\"Create Bedrock AgentCore Gateway\"\"\"\n",
    "    client = boto3.client(\n",
    "        'bedrock-agentcore-control',\n",
    "        region_name='us-east-1',\n",
    "        endpoint_url=\"https://bedrock-agentcore-control.us-east-1.amazonaws.com\"\n",
    "    )\n",
    "    \n",
    "    gateway_auth_config = {\n",
    "        \"customJWTAuthorizer\": {\n",
    "            \"allowedClients\": [auth_config['client_id']],\n",
    "            \"discoveryUrl\": auth_config['discovery_url']\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        response = client.create_gateway(\n",
    "            name=f'fruit-orders-dynamo-gateway-{int(time.time())%1000000}',\n",
    "            roleArn=role_arn,\n",
    "            protocolType='MCP',\n",
    "            authorizerType='CUSTOM_JWT',\n",
    "            authorizerConfiguration=gateway_auth_config,\n",
    "            description='Enhanced Bedrock Core Agent Gateway with Lambda integration'\n",
    "        )\n",
    "        \n",
    "        gateway_id = response.get('gatewayId')\n",
    "        print(f\"✅ Created Gateway: {gateway_id}\")\n",
    "        return gateway_id\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error creating gateway: {e}\")\n",
    "        return None\n",
    "\n",
    "if gateway_role_arn and auth_config:\n",
    "    gateway_id = create_bedrock_agentcore_gateway(gateway_role_arn, auth_config)\n",
    "    print(f\"Gateway ID: {gateway_id}\")\n",
    "else:\n",
    "    print(\"⚠️ Gateway role ARN and auth config required\")\n",
    "    gateway_id = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Create Gateway Lambda Target with Enhanced Schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_enhanced_gateway_target(gateway_id, lambda_arn):\n",
    "    client = boto3.client(\n",
    "        'bedrock-agentcore-control',\n",
    "        region_name='us-east-1',\n",
    "        endpoint_url=\"https://bedrock-agentcore-control.us-east-1.amazonaws.com\"\n",
    "    )\n",
    "    \n",
    "    lambda_target_config = {\n",
    "        \"mcp\": {\n",
    "            \"lambda\": {\n",
    "                \"lambdaArn\": lambda_arn,\n",
    "                \"toolSchema\": {\n",
    "                    \"inlinePayload\": [\n",
    "                        {\n",
    "                            \"name\": \"list_fruits_tool\",\n",
    "                            \"description\": \"Lists all available fruits with prices and stock levels\",\n",
    "                            \"inputSchema\": {\"type\": \"object\", \"properties\": {}, \"required\": []}\n",
    "                        },\n",
    "                        {\n",
    "                            \"name\": \"create_order_tool\",\n",
    "                            \"description\": \"Creates a new fruit order and updates inventory\",\n",
    "                            \"inputSchema\": {\n",
    "                                \"type\": \"object\",\n",
    "                                \"properties\": {\n",
    "                                    \"customer_name\": {\n",
    "                                        \"type\": \"string\", \n",
    "                                    },\n",
    "                                    \"items\": {\n",
    "                                        \"type\": \"array\",\n",
    "                                        \"items\": {\n",
    "                                            \"type\": \"object\",\n",
    "                                            \"properties\": {\n",
    "                                                \"name\": {\n",
    "                                                    \"type\": \"string\", \n",
    "                                                },\n",
    "                                                \"quantity\": {\n",
    "                                                    \"type\": \"number\", \n",
    "                                                }\n",
    "                                            },\n",
    "                                            \"required\": [\"name\", \"quantity\"]\n",
    "                                        }\n",
    "                                    }\n",
    "                                },\n",
    "                                \"required\": [\"customer_name\", \"items\"]\n",
    "                            }\n",
    "                        },\n",
    "                        {\n",
    "                            \"name\": \"get_order_tool\",\n",
    "                            \"description\": \"Gets order details by ID from DynamoDB\",\n",
    "                            \"inputSchema\": {\n",
    "                                \"type\": \"object\",\n",
    "                                \"properties\": {\n",
    "                                    \"order_id\": {\n",
    "                                        \"type\": \"string\",\n",
    "                                    }\n",
    "                                },\n",
    "                                \"required\": [\"order_id\"]\n",
    "                            }\n",
    "                        }\n",
    "                    ]\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        response = client.create_gateway_target(\n",
    "            gatewayIdentifier=gateway_id,\n",
    "            name='EnhancedFruitOrdersTarget',\n",
    "            description='Enhanced Lambda target with DynamoDB integration',\n",
    "            targetConfiguration=lambda_target_config,\n",
    "            credentialProviderConfigurations=[{\"credentialProviderType\": \"GATEWAY_IAM_ROLE\"}]\n",
    "        )\n",
    "        \n",
    "        target_id = response['targetId']\n",
    "        print(f\"✅ Created Target: {target_id}\")\n",
    "        return target_id\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error creating target: {e}\")\n",
    "        return None\n",
    "\n",
    "if gateway_id and enhanced_lambda_arn:\n",
    "    target_id = create_enhanced_gateway_target(gateway_id, enhanced_lambda_arn)\n",
    "    print(f\"Target ID: {target_id}\")\n",
    "else:\n",
    "    print(\"⚠️ Gateway ID and Lambda ARN required\")\n",
    "    target_id = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: Test Gateway with List and Ivoke "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_bearer_token(auth_config):\n",
    "    data = urllib.parse.urlencode({\n",
    "        \"grant_type\": \"client_credentials\",\n",
    "        \"client_id\": auth_config['client_id'],\n",
    "        \"client_secret\": auth_config['client_secret']\n",
    "    }).encode('utf-8')\n",
    "    \n",
    "    try:\n",
    "        req = urllib.request.Request(\n",
    "            auth_config['token_endpoint'],\n",
    "            data=data,\n",
    "            headers={\"Content-Type\": \"application/x-www-form-urlencoded\"}\n",
    "        )\n",
    "        \n",
    "        with urllib.request.urlopen(req, timeout=10) as response:\n",
    "            if response.status == 200:\n",
    "                result = json.loads(response.read().decode('utf-8'))\n",
    "                return result.get('access_token')\n",
    "    except Exception as e:\n",
    "        print(f\"Error generating token: {e}\")\n",
    "    return None\n",
    "\n",
    "def test_enhanced_gateway(gateway_id, target_id, auth_config):\n",
    "    bearer_token = generate_bearer_token(auth_config)\n",
    "    if not bearer_token:\n",
    "        print(\"❌ Failed to generate token\")\n",
    "        return\n",
    "    \n",
    "    mcp_endpoint = f\"https://{gateway_id}.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp\"\n",
    "    \n",
    "    # Test 1: List fruits with stock levels\n",
    "    print(\"\\n🧪 Test 1: List fruits with stock levels\")\n",
    "    mcp_request = {\n",
    "        \"jsonrpc\": \"2.0\",\n",
    "        \"id\": 1,\n",
    "        \"method\": \"tools/call\",\n",
    "        \"params\": {\n",
    "            \"name\": \"EnhancedFruitOrdersTarget___list_fruits_tool\",\n",
    "            \"arguments\": {}\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        req = urllib.request.Request(\n",
    "            mcp_endpoint,\n",
    "            data=json.dumps(mcp_request).encode('utf-8'),\n",
    "            headers={\n",
    "                'Content-Type': 'application/json',\n",
    "                'Authorization': f'Bearer {bearer_token}'\n",
    "            }\n",
    "        )\n",
    "        \n",
    "        with urllib.request.urlopen(req, timeout=30) as response:\n",
    "            if response.status == 200:\n",
    "                result = json.loads(response.read().decode('utf-8'))\n",
    "                print(\"✅ List fruits successful!\")\n",
    "                print(json.dumps(result, indent=2))\n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error: {e}\")\n",
    "    \n",
    "    # Test 2: Create order with inventory update\n",
    "    print(\"\\n🧪 Test 2: Create order with inventory update\")\n",
    "    create_order_request = {\n",
    "        \"jsonrpc\": \"2.0\",\n",
    "        \"id\": 2,\n",
    "        \"method\": \"tools/call\",\n",
    "        \"params\": {\n",
    "            \"name\": \"EnhancedFruitOrdersTarget___create_order_tool\",\n",
    "            \"arguments\": {\n",
    "                \"customer_name\": \"Alice Johnson\",\n",
    "                \"items\": [\n",
    "                    {\"name\": \"apple\", \"quantity\": 3},\n",
    "                    {\"name\": \"banana\", \"quantity\": 2}\n",
    "                ]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        req = urllib.request.Request(\n",
    "            mcp_endpoint,\n",
    "            data=json.dumps(create_order_request).encode('utf-8'),\n",
    "            headers={\n",
    "                'Content-Type': 'application/json',\n",
    "                'Authorization': f'Bearer {bearer_token}'\n",
    "            }\n",
    "        )\n",
    "        \n",
    "        with urllib.request.urlopen(req, timeout=30) as response:\n",
    "            if response.status == 200:\n",
    "                result = json.loads(response.read().decode('utf-8'))\n",
    "                print(\"✅ Create order successful!\")\n",
    "                print(json.dumps(result, indent=2))\n",
    "                \n",
    "                # Extract order ID for next test\n",
    "                if 'result' in result and 'content' in result['result']:\n",
    "                    content = result['result']['content'][0]['text']\n",
    "                    order_data = json.loads(content)\n",
    "                    if 'order' in order_data:\n",
    "                        order_id = order_data['order']['order_id']\n",
    "                        print(f\"\\n📝 Order ID for testing: {order_id}\")\n",
    "                        return order_id\n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error: {e}\")\n",
    "    \n",
    "    return None\n",
    "\n",
    "# Test the enhanced gateway\n",
    "if gateway_id and target_id and auth_config:\n",
    "    test_order_id = test_enhanced_gateway(gateway_id, target_id, auth_config)\n",
    "else:\n",
    "    print(\"⚠️ Gateway ID, Target ID, and auth config required for testing\")\n",
    "    test_order_id = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 9: Create Bedrock Agent that acts as MCP Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_enhanced_bedrock_agent():\n",
    "    bedrock = boto3.client('bedrock-agent', region_name='us-east-1')\n",
    "    \n",
    "    try:\n",
    "        agent_response = bedrock.create_agent(\n",
    "            agentName='enhanced-fruit-ordering-dynamo-agent',\n",
    "            description='Enhanced AI agent for fruit ordering with DynamoDB persistence',\n",
    "            foundationModel='anthropic.claude-3-sonnet-20240229-v1:0',\n",
    "            instruction='''You are an advanced fruit ordering assistant with access to a persistent inventory system. You can:\n",
    "\n",
    "1. **List Available Fruits**: Show all fruits with current prices and stock levels\n",
    "2. **Create Orders**: Process customer orders and automatically update inventory\n",
    "3. **Track Orders**: Retrieve order details using order IDs\n",
    "4. **Inventory Management**: Check stock levels and prevent overselling\n",
    "\n",
    "Key Features:\n",
    "- Real-time inventory tracking with DynamoDB\n",
    "- Automatic stock deduction when orders are placed\n",
    "- Persistent order storage with customer information\n",
    "- Stock validation to prevent overselling\n",
    "\n",
    "Always be helpful, provide clear information about availability, and confirm order details including total cost.''',\n",
    "            agentResourceRoleArn=gateway_role_arn,\n",
    "            idleSessionTTLInSeconds=1800\n",
    "        )\n",
    "        \n",
    "        agent_id = agent_response['agent']['agentId']\n",
    "        print(f\"✅ Created Enhanced Agent: {agent_id}\")\n",
    "        return agent_id\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error creating agent: {e}\")\n",
    "        return None\n",
    "\n",
    "enhanced_agent_id = create_enhanced_bedrock_agent()\n",
    "print(f\"Enhanced Agent ID: {enhanced_agent_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 10: Create Bridge Lambda for Bedrock Agent\n",
    "\n",
    "Create a Lambda function that bridges Bedrock Agent calls to the Bedrock AgentCore Gateway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_bridge_lambda(gateway_id, target_id, auth_config):\n",
    "    bridge_code = f'''\n",
    "import json\n",
    "import urllib.request\n",
    "import urllib.parse\n",
    "\n",
    "def generate_bearer_token():\n",
    "    data = urllib.parse.urlencode({{\n",
    "        \"grant_type\": \"client_credentials\",\n",
    "        \"client_id\": \"{auth_config['client_id']}\",\n",
    "        \"client_secret\": \"{auth_config['client_secret']}\"\n",
    "    }}).encode('utf-8')\n",
    "    \n",
    "    try:\n",
    "        req = urllib.request.Request(\n",
    "            \"{auth_config['token_endpoint']}\",\n",
    "            data=data,\n",
    "            headers={{\"Content-Type\": \"application/x-www-form-urlencoded\"}}\n",
    "        )\n",
    "        \n",
    "        with urllib.request.urlopen(req, timeout=10) as response:\n",
    "            if response.status == 200:\n",
    "                result = json.loads(response.read().decode('utf-8'))\n",
    "                return result.get('access_token')\n",
    "    except Exception as e:\n",
    "        print(f\"Error generating token: {{e}}\")\n",
    "    return None\n",
    "\n",
    "def lambda_handler(event, context):\n",
    "    gateway_endpoint = \"https://{gateway_id}.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp\"\n",
    "    \n",
    "    bearer_token = generate_bearer_token()\n",
    "    if not bearer_token:\n",
    "        return {{\n",
    "            \"messageVersion\": \"1.0\",\n",
    "            \"response\": {{\n",
    "                \"actionGroup\": event.get('actionGroup', ''),\n",
    "                \"function\": event.get('function', ''),\n",
    "                \"functionResponse\": {{\n",
    "                    \"responseBody\": {{\"TEXT\": {{\"body\": \"Failed to generate token\"}}}}\n",
    "                }}\n",
    "            }}\n",
    "        }}\n",
    "    \n",
    "    function_name = event.get('function')\n",
    "    parameters = event.get('parameters', [])\n",
    "    \n",
    "    param_dict = {{}}\n",
    "    for param in parameters:\n",
    "        param_dict[param['name']] = param['value']\n",
    "    \n",
    "    tool_mapping = {{\n",
    "        'list_fruits': 'EnhancedFruitOrdersTarget___list_fruits_tool',\n",
    "        'create_order': 'EnhancedFruitOrdersTarget___create_order_tool',\n",
    "        'get_order': 'EnhancedFruitOrdersTarget___get_order_tool'\n",
    "    }}\n",
    "    \n",
    "    gateway_tool = tool_mapping.get(function_name)\n",
    "    if not gateway_tool:\n",
    "        return {{\n",
    "            \"messageVersion\": \"1.0\",\n",
    "            \"response\": {{\n",
    "                \"actionGroup\": event.get('actionGroup', ''),\n",
    "                \"function\": function_name,\n",
    "                \"functionResponse\": {{\n",
    "                    \"responseBody\": {{\"TEXT\": {{\"body\": f\"Unknown function: {{function_name}}\"}}}}\n",
    "                }}\n",
    "            }}\n",
    "        }}\n",
    "    \n",
    "    arguments = {{}}\n",
    "    if function_name == 'create_order':\n",
    "        arguments = {{\n",
    "            'customer_name': param_dict.get('customer_name', 'Customer'),\n",
    "            'items': json.loads(param_dict.get('items', '[]'))\n",
    "        }}\n",
    "    elif function_name == 'get_order':\n",
    "        arguments = {{'order_id': param_dict.get('order_id')}}\n",
    "    \n",
    "    mcp_request = {{\n",
    "        \"jsonrpc\": \"2.0\",\n",
    "        \"id\": 1,\n",
    "        \"method\": \"tools/call\",\n",
    "        \"params\": {{\"name\": gateway_tool, \"arguments\": arguments}}\n",
    "    }}\n",
    "    \n",
    "    try:\n",
    "        req = urllib.request.Request(\n",
    "            gateway_endpoint,\n",
    "            data=json.dumps(mcp_request).encode('utf-8'),\n",
    "            headers={{\n",
    "                'Content-Type': 'application/json',\n",
    "                'Authorization': f'Bearer {{bearer_token}}'\n",
    "            }}\n",
    "        )\n",
    "        \n",
    "        with urllib.request.urlopen(req, timeout=30) as response:\n",
    "            if response.status == 200:\n",
    "                result = json.loads(response.read().decode('utf-8'))\n",
    "                if 'result' in result and 'content' in result['result']:\n",
    "                    content_text = result['result']['content'][0].get('text', '')\n",
    "                    try:\n",
    "                        parsed_result = json.loads(content_text)['response']['payload']\n",
    "                        response_text = json.dumps(parsed_result, indent=2)\n",
    "                    except:\n",
    "                        response_text = content_text\n",
    "                else:\n",
    "                    response_text = json.dumps(result, indent=2)\n",
    "            else:\n",
    "                response_text = f\"Gateway error: {{response.status}}\"\n",
    "    except Exception as e:\n",
    "        response_text = f\"Error: {{str(e)}}\"\n",
    "    \n",
    "    return {{\n",
    "        \"messageVersion\": \"1.0\",\n",
    "        \"response\": {{\n",
    "            \"actionGroup\": event.get('actionGroup', ''),\n",
    "            \"function\": function_name,\n",
    "            \"functionResponse\": {{\n",
    "                \"responseBody\": {{\"TEXT\": {{\"body\": response_text}}}}\n",
    "            }}\n",
    "        }}\n",
    "    }}\n",
    "'''\n",
    "    \n",
    "    lambda_client = boto3.client('lambda', region_name='us-east-1')\n",
    "    \n",
    "    with open('bridge_function.py', 'w') as f:\n",
    "        f.write(bridge_code)\n",
    "    \n",
    "    with zipfile.ZipFile('bridge_lambda.zip', 'w') as zip_file:\n",
    "        zip_file.write('bridge_function.py', 'lambda_function.py')\n",
    "    \n",
    "    try:\n",
    "        with open('bridge_lambda.zip', 'rb') as zip_file:\n",
    "            response = lambda_client.create_function(\n",
    "                FunctionName='enhanced-bedrock-gateway-bridge',\n",
    "                Runtime='python3.13',\n",
    "                Role=gateway_role_arn,\n",
    "                Handler='lambda_function.lambda_handler',\n",
    "                Code={'ZipFile': zip_file.read()},\n",
    "                Description='Bridge between Bedrock Agent and Bedrock AgentCore Gateway',\n",
    "                Timeout=30\n",
    "            )\n",
    "        \n",
    "        bridge_arn = response['FunctionArn']\n",
    "        print(f\"✅ Created Bridge Lambda: {bridge_arn}\")\n",
    "        \n",
    "        # Add permission for Bedrock\n",
    "        lambda_client.add_permission(\n",
    "            FunctionName='enhanced-bedrock-gateway-bridge',\n",
    "            StatementId='bedrock-agent-invoke',\n",
    "            Action='lambda:InvokeFunction',\n",
    "            Principal='bedrock.amazonaws.com'\n",
    "        )\n",
    "        \n",
    "        return bridge_arn\n",
    "        \n",
    "    except Exception as e:\n",
    "        if 'already exists' in str(e):\n",
    "            account_id = boto3.client('sts').get_caller_identity()['Account']\n",
    "            bridge_arn = f\"arn:aws:lambda:us-east-1:{account_id}:function:enhanced-bedrock-gateway-bridge\"\n",
    "            print(f\"✅ Bridge Lambda already exists: {bridge_arn}\")\n",
    "            return bridge_arn\n",
    "        else:\n",
    "            print(f\"❌ Error: {e}\")\n",
    "            return None\n",
    "    finally:\n",
    "        if os.path.exists('bridge_function.py'):\n",
    "            os.remove('bridge_function.py')\n",
    "        if os.path.exists('bridge_lambda.zip'):\n",
    "            os.remove('bridge_lambda.zip')\n",
    "\n",
    "# Create Bridge Lambda\n",
    "if gateway_id and target_id and auth_config:\n",
    "    bridge_arn = create_bridge_lambda(gateway_id, target_id, auth_config)\n",
    "    print(f\"Bridge Lambda ARN: {bridge_arn}\")\n",
    "else:\n",
    "    print(\"⚠️ Gateway ID, Target ID, and auth config required\")\n",
    "    bridge_arn = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 11: Wait and Verify Agent Creation\n",
    "\n",
    "Wait for the agent to be created and verify its status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wait_for_agent_creation(agent_id, max_wait=300):\n",
    "    \"\"\"Wait for agent to be created and ready\"\"\"\n",
    "    if not agent_id:\n",
    "        print(\"⚠️ No agent ID provided\")\n",
    "        return False\n",
    "        \n",
    "    bedrock = boto3.client('bedrock-agent', region_name='us-east-1')\n",
    "    \n",
    "    print(f\"⏳ Waiting for agent {agent_id} to be ready...\")\n",
    "    \n",
    "    start_time = time.time()\n",
    "    while time.time() - start_time < max_wait:\n",
    "        try:\n",
    "            response = bedrock.get_agent(agentId=agent_id)\n",
    "            status = response['agent']['agentStatus']\n",
    "            \n",
    "            print(f\"Agent status: {status}\")\n",
    "            \n",
    "            if status == 'NOT_PREPARED':\n",
    "                print(\"✅ Agent created successfully and ready for action groups\")\n",
    "                return True\n",
    "            elif status == 'FAILED':\n",
    "                print(\"❌ Agent creation failed\")\n",
    "                return False\n",
    "            \n",
    "            time.sleep(10)\n",
    "            \n",
    "        except Exception as e:\n",
    "            print(f\"Error checking agent status: {e}\")\n",
    "            time.sleep(10)\n",
    "    \n",
    "    print(\"⚠️ Timeout waiting for agent creation\")\n",
    "    return False\n",
    "\n",
    "# Wait for agent creation\n",
    "agent_ready = wait_for_agent_creation(enhanced_agent_id)\n",
    "print(f\"Agent ready: {agent_ready}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 12: Associate Action Group\n",
    "\n",
    "Create and associate the action group with the Bedrock Agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_action_group(agent_id, bridge_lambda_arn):\n",
    "    if not agent_id or not bridge_lambda_arn:\n",
    "        print(\"⚠️ Agent ID and Bridge Lambda ARN required\")\n",
    "        return None\n",
    "        \n",
    "    bedrock = boto3.client('bedrock-agent', region_name='us-east-1')\n",
    "    \n",
    "    try:\n",
    "        action_group_response = bedrock.create_agent_action_group(\n",
    "            agentId=agent_id,\n",
    "            agentVersion='DRAFT',\n",
    "            actionGroupName='EnhancedFruitOrdersActions',\n",
    "            description='Enhanced actions for fruit ordering with DynamoDB',\n",
    "            actionGroupExecutor={'lambda': bridge_lambda_arn},\n",
    "            functionSchema={\n",
    "                'functions': [\n",
    "                    {\n",
    "                        'name': 'list_fruits',\n",
    "                        'description': 'List available fruits with prices and stock levels'\n",
    "                    },\n",
    "                    {\n",
    "                        'name': 'create_order',\n",
    "                        'description': 'Create a new fruit order with inventory update',\n",
    "                        'parameters': {\n",
    "                            'customer_name': {'type': 'string', 'description': 'Customer name'},\n",
    "                            'items': {'type': 'string', 'description': 'JSON array of items with name and quantity'}\n",
    "                        }\n",
    "                    },\n",
    "                    {\n",
    "                        'name': 'get_order',\n",
    "                        'description': 'Get order details by ID from DynamoDB',\n",
    "                        'parameters': {\n",
    "                            'order_id': {'type': 'string', 'description': 'Order ID'}\n",
    "                        }\n",
    "                    }\n",
    "                ]\n",
    "            }\n",
    "        )\n",
    "        \n",
    "        action_group_id = action_group_response['agentActionGroup']['actionGroupId']\n",
    "        print(f\"✅ Created Action Group: {action_group_id}\")\n",
    "        \n",
    "        # Prepare agent\n",
    "        bedrock.prepare_agent(agentId=agent_id)\n",
    "        print(f\"✅ Agent prepared with action group\")\n",
    "        \n",
    "        return action_group_id\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error creating action group: {e}\")\n",
    "        return None\n",
    "\n",
    "# Create Action Group\n",
    "if agent_ready and bridge_arn:\n",
    "    action_group_id = create_action_group(enhanced_agent_id, bridge_arn)\n",
    "    print(f\"Action Group ID: {action_group_id}\")\n",
    "else:\n",
    "    print(\"⚠️ Agent must be ready and Bridge Lambda ARN required\")\n",
    "    action_group_id = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 13: Test Bedrock Agent End-to-End flow with Bedrock AgentCore Gateway with Lambda as target\n",
    "\n",
    "Test the complete integration by invoking the Bedrock Agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_enhanced_bedrock_agent(agent_id):\n",
    "    bedrock_runtime = boto3.client('bedrock-agent-runtime', region_name='us-east-1')\n",
    "    \n",
    "    # Check agent status first\n",
    "    bedrock = boto3.client('bedrock-agent', region_name='us-east-1')\n",
    "    try:\n",
    "        agent_info = bedrock.get_agent(agentId=agent_id)\n",
    "        agent_status = agent_info['agent']['agentStatus']\n",
    "        print(f\"🔍 Agent status: {agent_status}\")\n",
    "        \n",
    "        if agent_status != 'PREPARED':\n",
    "            print(f\"⚠️ Agent not prepared. Current status: {agent_status}\")\n",
    "            return\n",
    "    except Exception as e:\n",
    "        print(f\"❌ Error checking agent: {e}\")\n",
    "        return\n",
    "    \n",
    "    test_queries = [\n",
    "        \"What fruits are available?\",\n",
    "        \"Create an order for Bob with fruit as apple and quantity as 2\"\n",
    "    ]\n",
    "    \n",
    "    for query in test_queries:\n",
    "        print(f\"\\n🤔 User: {query}\")\n",
    "        \n",
    "        try:\n",
    "            response = bedrock_runtime.invoke_agent(\n",
    "                agentId=agent_id,\n",
    "                agentAliasId='TSTALIASID',\n",
    "                sessionId=f'test-session-{int(time.time())}',\n",
    "                inputText=query\n",
    "            )\n",
    "            \n",
    "            full_response = \"\"\n",
    "            for event in response['completion']:\n",
    "                if 'chunk' in event:\n",
    "                    chunk = event['chunk']\n",
    "                    if 'bytes' in chunk:\n",
    "                        full_response += chunk['bytes'].decode('utf-8')\n",
    "            \n",
    "            print(f\"🤖 Enhanced Agent: {full_response}\")\n",
    "            \n",
    "        except Exception as e:\n",
    "            print(f\"❌ Error: {e}\")\n",
    "\n",
    "# Test the Enhanced Bedrock Agent\n",
    "if enhanced_agent_id:\n",
    "    print(\"Testing Enhanced Bedrock Agent with DynamoDB integration...\")\n",
    "    test_enhanced_bedrock_agent(enhanced_agent_id)\n",
    "else:\n",
    "    print(\"⚠️ Agent ID required for testing\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🎉 Summary\n",
    "\n",
    "### Architecture:\n",
    "```\n",
    "User Query → Bedrock Agent → Bridge Lambda → Bedrock AgentCore Gateway → Enhanced Lambda → DynamoDB\n",
    "```\n",
    "\n",
    "### Database Schema:\n",
    "- **FruitOrders**: order_id, customer_name, items, total_cost, status, created_at\n",
    "- **FruitInventory**: fruit_name, price, unit, stock"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🧹 Cleanup Resources\n",
    "\n",
    "Run each section below to clean up specific resources. Execute in order to avoid dependency issues."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Delete Bedrock Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock = boto3.client('bedrock-agent', region_name='us-east-1')\n",
    "\n",
    "try:\n",
    "    if 'enhanced_agent_id' in globals() and enhanced_agent_id:\n",
    "        bedrock.delete_agent(agentId=enhanced_agent_id, skipResourceInUseCheck=True)\n",
    "        print(f\"✅ Deleted Enhanced Bedrock Agent: {enhanced_agent_id}\")\n",
    "    else:\n",
    "        print(\"⚠️ No enhanced_agent_id found\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ Enhanced Agent cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Delete Bridge Lambda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_client = boto3.client('lambda', region_name='us-east-1')\n",
    "\n",
    "try:\n",
    "    lambda_client.delete_function(FunctionName='enhanced-bedrock-gateway-bridge')\n",
    "    print(\"✅ Deleted Bridge Lambda\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ Bridge Lambda cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Delete Enhanced Lambda Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_client = boto3.client('lambda', region_name='us-east-1')\n",
    "\n",
    "try:\n",
    "    lambda_client.delete_function(FunctionName='enhanced-fruit-order-dynamo')\n",
    "    print(\"✅ Deleted Enhanced Lambda\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ Enhanced Lambda cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Delete Gateway Target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock_agentcore_client = boto3.client(\n",
    "    'bedrock-agentcore-control',\n",
    "    region_name='us-east-1',\n",
    "    endpoint_url=\"https://bedrock-agentcore-control.us-east-1.amazonaws.com\"\n",
    ")\n",
    "\n",
    "try:\n",
    "    if 'gateway_id' in globals() and gateway_id and 'target_id' in globals() and target_id:\n",
    "        bedrock_agentcore_client.delete_gateway_target(\n",
    "            gatewayIdentifier=gateway_id,\n",
    "            targetId=target_id\n",
    "        )\n",
    "        print(f\"✅ Deleted Gateway Target: {target_id}\")\n",
    "    else:\n",
    "        print(\"⚠️ No gateway_id or target_id found\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ Gateway Target cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Delete Bedrock AgentCore Gateway"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    if 'gateway_id' in globals() and gateway_id:\n",
    "        bedrock_agentcore_client.delete_gateway(gatewayIdentifier=gateway_id)\n",
    "        print(f\"✅ Deleted Bedrock AgentCore Gateway: {gateway_id}\")\n",
    "    else:\n",
    "        print(\"⚠️ No gateway_id found\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ Gateway cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6: Delete DynamoDB Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamodb = boto3.client('dynamodb', region_name='us-east-1')\n",
    "\n",
    "try:\n",
    "    dynamodb.delete_table(TableName='FruitOrders')\n",
    "    print(\"✅ Deleted FruitOrders table\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ FruitOrders table cleanup error: {e}\")\n",
    "\n",
    "try:\n",
    "    dynamodb.delete_table(TableName='FruitInventory')\n",
    "    print(\"✅ Deleted FruitInventory table\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ FruitInventory table cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7: Delete Cognito Resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cognito = boto3.client('cognito-idp', region_name='us-east-1')\n",
    "\n",
    "try:\n",
    "    if 'auth_config' in globals() and auth_config:\n",
    "        user_pool_id = auth_config.get('user_pool_id')\n",
    "        domain_name = auth_config.get('domain_name')\n",
    "        \n",
    "        if user_pool_id and user_pool_id != 'us-east-1_0IalqBD2P':\n",
    "            # Delete domain first\n",
    "            try:\n",
    "                cognito.delete_user_pool_domain(\n",
    "                    Domain=domain_name,\n",
    "                    UserPoolId=user_pool_id\n",
    "                )\n",
    "                print(f\"✅ Deleted Cognito Domain: {domain_name}\")\n",
    "            except Exception as domain_error:\n",
    "                print(f\"⚠️ Domain deletion error: {domain_error}\")\n",
    "            \n",
    "            # Delete user pool\n",
    "            cognito.delete_user_pool(UserPoolId=user_pool_id)\n",
    "            print(f\"✅ Deleted Cognito User Pool: {user_pool_id}\")\n",
    "        else:\n",
    "            print(\"⚠️ Using default user pool, skipping deletion\")\n",
    "    else:\n",
    "        print(\"⚠️ No auth_config found\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️ Cognito cleanup error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 8: Delete IAM Role"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam = boto3.client('iam')\n",
    "\n",
    "try:\n",
    "    # Detach policies first\n",
    "    iam.detach_role_policy(\n",
    "        RoleName='BedrockAgentCoreGatewayLambdaRole',\n",
    "        PolicyArn='arn:aws:iam::aws:policy/AdministratorAccess'\n",
    "    )\n",
    "    print(\"✅ Detached AdministratorAccess policy\")\n",
    "    iam.delete_role_policy(\n",
    "        RoleName='BedrockAgentCoreGatewayLambdaRole',\n",
    "        PolicyName='BedrockAgentCoreInlinePolicy'\n",
    "    )\n",
    "    # Delete role\n",
    "    iam.delete_role(RoleName='BedrockAgentCoreGatewayLambdaRole')\n",
    "    print(\"✅ Deleted IAM Role: BedrockAgentCoreGatewayLambdaRole\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"⚠️ IAM Role cleanup error: {e}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
